import math as pm
import numpy as np
from numpy import genfromtxt
from scipy import stats

PREFIX_U = 'u_r_'
PREFIX_V = 'v_r_'
PREFIX_RHO = 'rho_r_'
PATH = '/home/behollis/DATA/in/ts00050/'
DPATH = '/home/behollis/DATA/out/lockExSt/ts00050/'

gUv = dict()   
DIM = 2

prob = None
slines = None

'''
Reads velocity fields. Calculates KDE for each interpolated velocity PDF. Assigns
velocity probability density to displacement (edge) in streamline curve.
'''

def lerp(a, b, w):
    return a + (b - a) * w

def readMember(mem):
    ''' Read numpy matrix from disk and return tuple of realization data. '''
    filename_u = PREFIX_U + str(mem) + '.txt'
    filename_v = PREFIX_V + str(mem) + '.txt'
    #filename_rho = PREFIX_RHO + str(mem) + '.txt'
    
    u = genfromtxt(PATH + filename_u)
    u = -1.0*u[0:127,0:127]
    
    v = genfromtxt(PATH + filename_v)
    #rho = genfromtxt(PATH + filename_rho)
    
    return (u, v)# , rho)

def getCoordParts(ppos=[0.0,0.0]):
    #decompose fract / whole from particle position
    ppos_parts = [[0.0,0.0],[0.0,0.0]] #[fract,whole] for each x,y comp
    ppos_parts[0][0] = pm.modf(ppos[0])[0];ppos_parts[0][1] = pm.modf(ppos[0])[1]
    ppos_parts[1][0] = pm.modf(ppos[1])[0];ppos_parts[1][1] = pm.modf(ppos[1])[1]
    
    return ppos_parts 

def getGridPoints(ppos=[0.0,0.0]):
    #assume grid points are defined by integer indices
    
    ppos_parts = getCoordParts(ppos)
    
    #print "quantile alpha x: " + str( ppos_parts[0][0] )
    #print "quantile alpha y: " + str( ppos_parts[1][0] )
    
    # grid point numbers:
    #
    # (2)---(3)
    # |      |
    # |      |
    # (0)---(1)
    
    #find four corner grid point indices, numbered from gpt0 = (bottom, left) TO gpt3 = (top, right)
    #calculated from whole parts 
    gpt0 = [ppos_parts[0][1], ppos_parts[1][1]]
    gpt1 = [ppos_parts[0][1] + 1, ppos_parts[1][1]]
    gpt2 = [ppos_parts[0][1], ppos_parts[1][1] + 1]
    gpt3 = [ppos_parts[0][1] + 1, ppos_parts[1][1] + 1]
    
    return gpt0, gpt1, gpt2, gpt3

def getVclinSamples(gpt0, gpt1, gpt2, gpt3, MEM):
    
    MEMS = len( gUv.keys() )
    
    gpt0_dist = np.zeros(shape=(DIM,MEMS))
    gpt1_dist = np.zeros(shape=(DIM,MEMS))
    gpt2_dist = np.zeros(shape=(DIM,MEMS))
    gpt3_dist = np.zeros(shape=(DIM,MEMS))
    
    #for idx in range(1,MEM):
    idx = 0
    for mem in gUv.keys():
        midx = mem#int(mem[3:])
        try:
            gpt0_dist[0][idx] = gUv[midx][0][gpt0[0]][gpt0[1]]
            gpt0_dist[1][idx] = gUv[midx][1][gpt0[0]][gpt0[1]]
            
            gpt1_dist[0][idx] = gUv[midx][0][gpt1[0]][gpt1[1]]
            gpt1_dist[1][idx] = gUv[midx][1][gpt1[0]][gpt1[1]]
            
            gpt2_dist[0][idx] = gUv[midx][0][gpt2[0]][gpt2[1]]
            gpt2_dist[1][idx] = gUv[midx][1][gpt2[0]][gpt2[1]]
            
            gpt3_dist[0][idx] = gUv[midx][0][gpt3[0]][gpt3[1]]
            gpt3_dist[1][idx] = gUv[midx][1][gpt3[0]][gpt3[1]]
            
            idx += 1
        except:
            return None
   
    return gpt0_dist, gpt1_dist, gpt2_dist, gpt3_dist


def interpVelFromEnsemble(ppos, MEM):
    
    ppos_parts = getCoordParts(ppos)
    
    #find grid cell that contains position to be interpolated
    gpt0, gpt1, gpt2, gpt3 = getGridPoints(ppos)
    
    gpt0_dist, gpt1_dist, gpt2_dist, gpt3_dist = getVclinSamples(gpt0, gpt1, gpt2, gpt3, MEM)
    
    #lerp ensemble samples
    try:
        lerp_u_gp0_gp1 = lerp( np.asarray( gpt0_dist[0] ), np.asarray( gpt1_dist[0]), w = ppos_parts[0][0] )
        lerp_u_gp2_gp3 = lerp( np.asarray( gpt2_dist[0] ), np.asarray( gpt3_dist[0]), w = ppos_parts[0][0] ) 
        lerp_u = lerp( np.asarray(lerp_u_gp0_gp1), np.asarray(lerp_u_gp2_gp3), w = ppos_parts[1][0] )  
        
        lerp_v_gp0_gp1 = lerp( np.asarray(gpt0_dist[1] ), np.asarray(gpt1_dist[1]), w = ppos_parts[0][0] )
        lerp_v_gp2_gp3 = lerp( np.asarray(gpt2_dist[1] ), np.asarray(gpt3_dist[1]), w = ppos_parts[0][0] ) 
        lerp_v = lerp( np.asarray(lerp_v_gp0_gp1), np.asarray(lerp_v_gp2_gp3), w = ppos_parts[1][0] )  
            
        k = stats.kde.gaussian_kde((lerp_u,lerp_v)) 
        return k, lerp_u.min(axis=0), lerp_u.max(axis=0), lerp_v.min(axis=0), lerp_v.max(axis=0)
    except:
        print 'streamlines off grid'
        return None

def calcStreamlineProbs(seed_pt, slines, MEM):
    slines_prob = dict()
    
    for m in range(1, MEM):
        gUv[m] = readMember(m) 
    
    probs_list = list()
    
#    for slidx, sl in enumerate(slines):
    for slidx in range(1, MEM-5):
        sl = slines[slidx]
        tot = 0
        slines_prob[slidx] = [ None, [] ]
        for ptidx in range(0, len(sl[0])):
            
            pt = [ sl[0][ptidx], sl[1][ptidx] ]
            
            kde, min_u, max_u, min_v, max_v = interpVelFromEnsemble(pt[0:2], MEM)
            
            #plotDistro( kde, (min_u,max_u), (min_v,max_v), '') 
            #band_width = MEM**(-1./(DIM+4))
            epsx = 0.05 * pm.fabs(max_u - min_u)
            epsy = 0.05 * pm.fabs(max_v - min_v)
              
            m = slidx + 1# member and slidx are synonymous
            
            u = gUv[m][0][pt[0],pt[1]]
            v = gUv[m][1][pt[0],pt[1]]
            low_bounds = [u - epsx, v - epsy]
            high_bounds = [u + epsx, v + epsy]
            prob = kde.integrate_box( low_bounds, high_bounds )
            tot += prob 
            slines_prob[slidx][1].append(prob)
            
        probs_list.append(tot)
            
                
        print('streamline {0}: {1}'.format(slidx, tot))
        slines_prob[slidx][0] = tot
        
        
    return slines_prob, probs_list
